import os
import json

def IMG(path):
    EXT = ('.jpg', '.jpeg', '.png', '.tif', '.tiff')
    for ext in EXT:
        _ext = os.path.splitext(path)[1]
        if _ext in (ext, ext.upper()):
            return True
    return False

def load_data(src):
    data = []
    for root, dirs, files in os.walk(src):
        for file in files:
            imgpath = os.path.join(root, file)
            if not IMG(imgpath):
                continue
            annpath = os.path.splitext(imgpath)[0] + '.json'
            if not os.path.exists(annpath):
                continue
            data.append([imgpath, annpath])
    return data

def load_ann(path):
    with open(path) as fp:
        return json.load(fp)

def write_ann(ann_info, path):
    with open(path, 'w') as fp:
        json.dump(ann_info, fp)

def instance2semantic(src, dst):
    data_infos = load_data(src)
    if not os.path.exists(dst):
        os.makedirs(dst)

    n = len(data_infos)
    for i, (imgsrc, annsrc) in enumerate(data_infos,1):
        print('[%d/%d] %s' % (i, n, os.path.basename(imgsrc)))
        ann_info = load_ann(annsrc)
        shapes = ann_info['shapes']
        temp = []
        for shape in shapes:
            label = shape['label']
            tags = label.split('-')
            if len(tags) and tags[-1].isdigit():
                tag = '-%s' % tags[-1]
                label = label[:-len(tag)]
            shape['label'] = label
            temp += [shape]
        ann_info['shapes'] = temp
        ann_info['imagePath'] = os.path.basename(imgsrc)
        ann_info['imageData'] = None
        write_ann(ann_info, os.path.join(dst, os.path.basename(annsrc)))


# src = r'G:\data\datasets\drink\rio\ai\train'
# main(src, src)
